# app/routes/auth.py
import jwt
from flask import Blueprint, request, jsonify, current_app
from datetime import datetime, timedelta, timezone # 确保 timezone 被导入
from app import db, bcrypt
from app.models import User
from functools import wraps

bp = Blueprint('auth', __name__)

@bp.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    if not data or not data.get('username') or not data.get('password'):
        return jsonify({"error": "必须提供用户名和密码"}), 400

    if User.query.filter_by(username=data['username']).first():
        return jsonify({"error": "用户名已存在"}), 400

    # 默认角色是 'normal' (在模型中设置)
    role = data.get('role', 'normal')
    if role not in ['normal', 'admin']: # 基本验证
         return jsonify({"error": "无效的角色"}), 400

    user = User(username=data['username'], role=role)
    user.set_password(data['password']) # 哈希密码
    db.session.add(user)
    db.session.commit()

    return jsonify({"message": "用户注册成功"}), 201

@bp.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    if not data or not data.get('username') or not data.get('password'):
        return jsonify({"error": "必须提供用户名和密码"}), 400

    user = User.query.filter_by(username=data['username']).first()

    if user and user.check_password(data['password']):
        # 生成 JWT Token
        payload = {
            'user_id': user.id,
            'username': user.username,
            'role': user.role,
            # 使用 UTC 时间并设置过期时间
            'exp': datetime.now(timezone.utc) + timedelta(hours=24) # Token 24 小时后过期
        }
        token = jwt.encode(payload, current_app.config['SECRET_KEY'], algorithm='HS256')

        return jsonify({
            "message": "登录成功",
            "token": token,
            "user": { # 返回一些基本用户信息给前端
                "id": user.id,
                "username": user.username,
                "role": user.role
            }
        }), 200
    else:
        # 不要提示是用户名错了还是密码错了
        return jsonify({"error": "无效的用户名或密码"}), 401

# JWT 验证装饰器
def token_required(f=None, *, required_role=None):
    if f is None:
        # Allows passing arguments like @token_required(required_role='admin')
        from functools import partial
        return partial(token_required, required_role=required_role)

    @wraps(f)
    def decorated_function(*args, **kwargs):
        token = None
        # 检查请求头中是否有 'Authorization'
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            try:
                # 提取 'Bearer <token>' 中的 token 部分
                token = auth_header.split(" ")[1]
            except IndexError:
                return jsonify({"error": "无效的 Authorization header 格式"}), 401

        if not token:
            return jsonify({"error": "缺少 Token"}), 401

        try:
            # 使用应用的 SECRET_KEY 解码和验证 token
            data = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=["HS256"])
            # 从 token 中获取 user_id，并查询数据库获取用户对象
            current_user = User.query.get(data['user_id'])
            if not current_user:
                 return jsonify({"error": "找不到用户"}), 401 # Token 有效但用户已被删除

            # 检查角色权限 (如果装饰器指定了 required_role)
            if required_role and current_user.role != required_role:
                 return jsonify({"error": "权限不足"}), 403 # Forbidden

        except jwt.ExpiredSignatureError:
            return jsonify({"error": "Token 已过期"}), 401
        except jwt.InvalidTokenError:
            return jsonify({"error": "无效的 Token"}), 401
        except Exception as e:
            # 捕获其他可能的解码错误
            current_app.logger.error(f"Token validation error: {e}") # 记录日志
            return jsonify({"error": "Token 验证失败"}), 401

        # 将获取到的用户对象注入到路由函数中
        # 注意：我们通常将 current_user 作为第一个参数传递
        return f(current_user, *args, **kwargs)

    return decorated_function

# (可以添加一个 /me 路由来测试装饰器)
@bp.route('/me', methods=['GET'])
@token_required
def get_me(current_user):
    """获取当前登录用户的信息"""
    return jsonify({
        "id": current_user.id,
        "username": current_user.username,
        "role": current_user.role
    }), 200


